"""
添加plan_date和batch_number字段到schedule_results和time_slots表
"""

import os
import sys
import logging
from datetime import datetime

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 导入Flask应用
from app import create_app, db
from sqlalchemy import text

def add_columns():
    """添加plan_date和batch_number字段到schedule_results和time_slots表"""
    logger.info("开始添加plan_date和batch_number字段...")

    # 创建Flask应用
    app = create_app()

    # 在应用上下文中执行操作
    with app.app_context():
        try:
            # 检查schedule_results表中的字段是否已存在
            result = db.session.execute(text("SHOW COLUMNS FROM schedule_results LIKE 'plan_date'"))
            if result.rowcount == 0:
                # 添加plan_date字段
                db.session.execute(text("ALTER TABLE schedule_results ADD COLUMN plan_date DATE NULL"))
                logger.info("成功添加plan_date字段到schedule_results表")
            else:
                logger.info("plan_date字段已存在于schedule_results表，无需添加")

            result = db.session.execute(text("SHOW COLUMNS FROM schedule_results LIKE 'batch_number'"))
            if result.rowcount == 0:
                # 添加batch_number字段
                db.session.execute(text("ALTER TABLE schedule_results ADD COLUMN batch_number INT NULL"))
                logger.info("成功添加batch_number字段到schedule_results表")
            else:
                logger.info("batch_number字段已存在于schedule_results表，无需添加")

            # 检查time_slots表中的字段是否已存在
            result = db.session.execute(text("SHOW COLUMNS FROM time_slots LIKE 'plan_date'"))
            if result.rowcount == 0:
                # 添加plan_date字段
                db.session.execute(text("ALTER TABLE time_slots ADD COLUMN plan_date DATE NULL"))
                logger.info("成功添加plan_date字段到time_slots表")
            else:
                logger.info("plan_date字段已存在于time_slots表，无需添加")

            result = db.session.execute(text("SHOW COLUMNS FROM time_slots LIKE 'batch_number'"))
            if result.rowcount == 0:
                # 添加batch_number字段
                db.session.execute(text("ALTER TABLE time_slots ADD COLUMN batch_number INT NULL"))
                logger.info("成功添加batch_number字段到time_slots表")
            else:
                logger.info("batch_number字段已存在于time_slots表，无需添加")

            # 添加create_time字段
            result = db.session.execute(text("SHOW COLUMNS FROM schedule_results LIKE 'create_time'"))
            if result.rowcount == 0:
                db.session.execute(text("ALTER TABLE schedule_results ADD COLUMN create_time DATETIME NULL DEFAULT CURRENT_TIMESTAMP"))
                logger.info("成功添加create_time字段到schedule_results表")
            else:
                logger.info("create_time字段已存在于schedule_results表，无需添加")

            result = db.session.execute(text("SHOW COLUMNS FROM time_slots LIKE 'create_time'"))
            if result.rowcount == 0:
                db.session.execute(text("ALTER TABLE time_slots ADD COLUMN create_time DATETIME NULL DEFAULT CURRENT_TIMESTAMP"))
                logger.info("成功添加create_time字段到time_slots表")
            else:
                logger.info("create_time字段已存在于time_slots表，无需添加")

            # 添加update_time字段
            result = db.session.execute(text("SHOW COLUMNS FROM schedule_results LIKE 'update_time'"))
            if result.rowcount == 0:
                db.session.execute(text("ALTER TABLE schedule_results ADD COLUMN update_time DATETIME NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP"))
                logger.info("成功添加update_time字段到schedule_results表")
            else:
                logger.info("update_time字段已存在于schedule_results表，无需添加")

            result = db.session.execute(text("SHOW COLUMNS FROM time_slots LIKE 'update_time'"))
            if result.rowcount == 0:
                db.session.execute(text("ALTER TABLE time_slots ADD COLUMN update_time DATETIME NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP"))
                logger.info("成功添加update_time字段到time_slots表")
            else:
                logger.info("update_time字段已存在于time_slots表，无需添加")

            # 添加status字段
            result = db.session.execute(text("SHOW COLUMNS FROM schedule_results LIKE 'status'"))
            if result.rowcount == 0:
                db.session.execute(text("ALTER TABLE schedule_results ADD COLUMN status VARCHAR(20) NULL DEFAULT 'active'"))
                logger.info("成功添加status字段到schedule_results表")
            else:
                logger.info("status字段已存在于schedule_results表，无需添加")

            result = db.session.execute(text("SHOW COLUMNS FROM time_slots LIKE 'status'"))
            if result.rowcount == 0:
                db.session.execute(text("ALTER TABLE time_slots ADD COLUMN status VARCHAR(20) NULL DEFAULT 'active'"))
                logger.info("成功添加status字段到time_slots表")
            else:
                logger.info("status字段已存在于time_slots表，无需添加")

            # 提交事务
            db.session.commit()

            # 添加索引
            try:
                db.session.execute(text("CREATE INDEX idx_schedule_results_plan_date ON schedule_results(plan_date)"))
                logger.info("成功添加plan_date索引到schedule_results表")
            except Exception as e:
                logger.warning(f"添加plan_date索引到schedule_results表失败，可能已存在: {str(e)}")

            try:
                db.session.execute(text("CREATE INDEX idx_schedule_results_batch_number ON schedule_results(batch_number)"))
                logger.info("成功添加batch_number索引到schedule_results表")
            except Exception as e:
                logger.warning(f"添加batch_number索引到schedule_results表失败，可能已存在: {str(e)}")

            try:
                db.session.execute(text("CREATE INDEX idx_time_slots_plan_date ON time_slots(plan_date)"))
                logger.info("成功添加plan_date索引到time_slots表")
            except Exception as e:
                logger.warning(f"添加plan_date索引到time_slots表失败，可能已存在: {str(e)}")

            try:
                db.session.execute(text("CREATE INDEX idx_time_slots_batch_number ON time_slots(batch_number)"))
                logger.info("成功添加batch_number索引到time_slots表")
            except Exception as e:
                logger.warning(f"添加batch_number索引到time_slots表失败，可能已存在: {str(e)}")

            db.session.commit()

            return True
        except Exception as e:
            db.session.rollback()
            logger.error(f"添加字段失败: {str(e)}")
            return False

if __name__ == "__main__":
    # 添加字段
    success = add_columns()

    if success:
        print("成功添加plan_date和batch_number字段到schedule_results和time_slots表")
    else:
        print("添加字段失败")
        sys.exit(1)
